package com.apps.sdses.flink141.source;

import com.clickhouse.client.internal.google.common.collect.Maps;
import lombok.extern.slf4j.Slf4j;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;
import ru.yandex.clickhouse.ClickHouseConnection;
import ru.yandex.clickhouse.ClickHouseDataSource;
import ru.yandex.clickhouse.ClickHouseStatement;
import ru.yandex.clickhouse.settings.ClickHouseProperties;
import ru.yandex.clickhouse.settings.ClickHouseQueryParam;

import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.time.LocalDate;
import java.time.format.DateTimeFormatter;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;

@Slf4j
public class ClickhouseSource extends RichParallelSourceFunction<Map<String, String>> {

    private boolean flag = true;
    private ClickHouseConnection conn = null;
    private ClickHouseStatement stmt = null;
    private ResultSet rs = null;
    private Map<ClickHouseQueryParam, String> additionalDBParams = new HashMap<>();

    // open只执行一次,适合开启资源
    @Override
    public void open(Configuration parameters) throws Exception {
        String url = "jdbc:clickhouse://172.18.133.12:8123";

        ClickHouseProperties properties = new ClickHouseProperties();
        properties.setSessionId(UUID.randomUUID().toString());
        properties.setUser("default");
        properties.setPassword("Deer@20161010");
        properties.setDatabase("CB6_0_ATOM");

        ClickHouseDataSource dataSource = new ClickHouseDataSource(url, properties);
        // ClickHouseProperties
        additionalDBParams.put(ClickHouseQueryParam.SESSION_ID, UUID.randomUUID().toString());
        conn = dataSource.getConnection();
        stmt = conn.createStatement();
    }

    @Override
    public void run(SourceContext<Map<String, String>> ctx) throws Exception {

        String statement = "select WDY_ID, RealTime, SaveTime, TaskTime, ID, PD_ID, MpNo, EquipNo, EquipTypeCode, WdyTEMP, WdyHUM, AlarmStatus, ErrorDetail, ExtendField, Status" +
                " from Equip_WdyHisData where SaveTime between '%s 00:00:00' and '%s 23:59:59' limit %s,%s";
        String start = "2023-11-15";
        String end = "2023-11-15";
//        String end = "2024-03-15";
        int limitStart = 0;
        int limit = 50000;
        while (flag) {
            String sql = String.format(statement, start, end, limitStart, limit);

            rs = stmt.executeQuery(sql, additionalDBParams);
            int c = 0;
            while (rs.next()) {
                c++;
                Map<String, String> hm = Maps.newHashMap();
                ResultSetMetaData rsmd = rs.getMetaData();
                int count = rsmd.getColumnCount();// 获取列的数量
                for (int i = 1; i <= count; i++) {
                    String key = rsmd.getColumnLabel(i);

                    //对指定字段名进行转换
                    if ("WDY_ID".equals(key)) {
                        key = "wdy_id";
                    } else if ("RealTime".equals(key)) {
                        key = "real_time";
                    } else if ("SaveTime".equals(key)) {
                        key = "save_time";
                    } else if ("TaskTime".equals(key)) {
                        key = "task_time";
                    } else if ("ID".equals(key)) {
                        key = "id";
                    } else if ("PD_ID".equals(key)) {
                        key = "pd_id";
                    } else if ("MpNo".equals(key)) {
                        key = "mpno";
                    } else if ("EquipNo".equals(key)) {
                        key = "equip_no";
                    } else if ("EquipTypeCode".equals(key)) {
                        key = "equip_type_code";
                    } else if ("WdyTEMP".equals(key)) {
                        key = "wdy_temp";
                    } else if ("WdyHUM".equals(key)) {
                        key = "wdy_hum";
                    } else if ("AlarmStatus".equals(key)) {
                        key = "alarm_status";
                    } else if ("ErrorDetail".equals(key)) {
                        key = "error_detail";
                    } else if ("ExtendField".equals(key)) {
                        key = "extend_field";
                    } else if ("Status".equals(key)) {
                        key = "status";
                    }
                    String value = rs.getString(i);
                    hm.put(key, value);
                }
                ctx.collect(hm);
            }

            limitStart += c;

            if (c < limit) {
                // 查完了，换时间
                limitStart = 0;
                LocalDate localDate = LocalDate.parse(start, DateTimeFormatter.ofPattern("yyyy-MM-dd")).plusDays(1);
                start = localDate.toString();
                end = localDate.toString();

                if (localDate.isAfter(LocalDate.now())) {
                    break;
                }
            }
        }
    }

    // 接收到cancel命令时取消数据生成
    @Override
    public void cancel() {
        flag = false;
    }

    @Override
    public void close() throws Exception {
        if (conn != null)
            conn.close();
        if (stmt != null)
            stmt.close();
        if (rs != null)
            rs.close();
    }
}
